package com.cloudea.learn.example.function;


import com.cloudea.learn.Model;
import com.cloudea.learn.Tensor;
import com.cloudea.learn.component.Component;
import com.cloudea.learn.layer.*;
import com.cloudea.learn.component.Linear;
import com.cloudea.learn.metric.MeanSquareError;
import com.cloudea.learn.optimizer.Adam;

import java.util.List;
import java.util.LinkedList;
import java.util.Random;

class MLP extends Component {

    private final Component linear1 = new Linear(3, 5);
    private final Component linear2 = new Linear(5, 7);
    private final Component linear3 = new Linear(7, 3);
    private final Layer relu = new Elu();

    @Override
    public Tensor forward(Tensor... x) {
        Tensor input = x[0];
        input = linear1.forward(input);
        input = relu.call(input);
        input = linear2.forward(input);
        input = relu.call(input);
        input = linear3.forward(input);
        input = relu.call(input);
        return input;
    }
}


public class main {

    static double f1(double x, double y, double z){
        return Math.sqrt(x*x+y*y+z*z);
    }
    static double f2(double x, double y, double z){
        return x + y+ z ;
    }
    static double f3(double x, double y, double z){
        return x * y * z;
    }

    public static void main(String[] args) {

        Model model = new Model(new MLP());

        int epoch = 64;
        int batch = 64;
        Random random = new Random();
        for(int i = 0; i < epoch; i++){
            System.out.println( String.format("########################################### epoch: %s", i + 1));
            List<Tensor[]> X = new LinkedList<>();
            List<Tensor> Y = new LinkedList<>();
            for(int j = 0; j < batch; j++){
                Tensor x = new Tensor(3);
                Tensor y = new Tensor(3);
                double x1 = random.nextDouble();
                double x2 = random.nextDouble();
                double x3 = random.nextDouble();
                x.set(x1, 0);
                x.set(x2, 1);
                x.set(x3, 2);
                y.set(f1(x1, x2, x3), 0);
                y.set(f2(x1, x2, x3), 1);
                y.set(f3(x1, x2, x3), 2);
                X.add(new Tensor[]{x});
                Y.add(y);
            }
            for (int j = 0; j < 64; j++){
                model.train(X, Y, new MeanSquareError(), new Adam.Builder());
            }
        }

        //预测
        double x1  = 1;
        double x2 = 0.2;
        double x3 = 0.3;
        double y1 = f1(x1, x2, x3);
        double y2 = f2(x1, x2, x3);
        double y3 = f3(x1, x2, x3);
        Tensor x = new Tensor(3);
        x.set(x1, 0);
        x.set(x2, 1);
        Tensor predict = model.predict(new Tensor[]{x});
        System.out.println(String.format("truth: %s, %s, %s", y1, y2, y3));
        System.out.println(predict);
    }
}
